{
 "cells": [
  {
   "cell_type": "markdown",
   "metadata": {},
   "source": [
    "https://mp.weixin.qq.com/s/IjpGTuMGBYeNaxNawmj6Yg"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 17,
   "metadata": {},
   "outputs": [
    {
     "name": "stdout",
     "output_type": "stream",
     "text": [
      "Loading bar data...\n",
      "Loaded bar data: 0:00:00 \n",
      "\n",
      "Loading bar data...\n",
      "Loaded bar data: 0:00:00 \n",
      "\n",
      "Loading bar data...\n",
      "Loaded bar data: 0:00:00 \n",
      "\n"
     ]
    }
   ],
   "source": [
    "from pybroker.ext.data import AKShare\n",
    "import pandas as pd\n",
    "import torch\n",
    "import torch.nn as nn\n",
    "import matplotlib.pyplot as plt\n",
    "\n",
    "ak = AKShare()\n",
    "symbols = [\"000001\", \"600159\", \"600733\"]\n",
    "\n",
    "# Descargar los datos de precios\n",
    "prices = [ak.query(symbols=symbol, start_date=\"2019-01-01\", end_date=\"2024-01-01\", adjust=\"qfq\")['close'] for symbol in symbols]\n",
    "data = pd.DataFrame(prices).T\n",
    "data.columns=symbols\n",
    "data = data.dropna() # 删除缺失值\n",
    "# 转换为 numpy，然后转换为 torch 张量\n",
    "prices = data.values # 形状：（天数， 股票数量）\n",
    "prices_tensor = torch.tensor(prices, dtype=torch.float32) # 转换为 PyTorch 张量"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 18,
   "metadata": {},
   "outputs": [],
   "source": [
    "class FourierConvolutionModel(nn.Module):\n",
    "  def __init__(self):\n",
    "\n",
    "    super(FourierConvolutionModel, self).__init__()\n",
    "    self.conv1 = nn.Conv1d(in_channels=10, out_channels=10, kernel_size=3, padding=1)\n",
    "    self.relu = nn.ReLU()\n",
    "    self.fc = nn.Linear(10 * 60, 5) # Fully connected layer to predict prices for 5 stocks\n",
    "\n",
    "  def forward(self, x):\n",
    "    # Combine real and imaginary parts along the channel dimension\n",
    "    x = torch.cat([x.real, x.imag], dim=1)\n",
    "    x = self.conv1(x) # Convolution in frequency domain\n",
    "    x = fft.ifft(x, dim=2).real # Inverse Fourier Transform on the real part\n",
    "    x = x.view(x.size(0), -1)\n",
    "    x = self.relu(x)\n",
    "    x = self.fc(x)\n",
    "    return x"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": 19,
   "metadata": {},
   "outputs": [
    {
     "ename": "AttributeError",
     "evalue": "'Tensor' object has no attribute 'append'",
     "output_type": "error",
     "traceback": [
      "\u001b[1;31m---------------------------------------------------------------------------\u001b[0m",
      "\u001b[1;31mAttributeError\u001b[0m                            Traceback (most recent call last)",
      "Cell \u001b[1;32mIn[19], line 6\u001b[0m\n\u001b[0;32m      3\u001b[0m X_train, y_train \u001b[38;5;241m=\u001b[39m [], []\n\u001b[0;32m      5\u001b[0m \u001b[38;5;28;01mfor\u001b[39;00m i \u001b[38;5;129;01min\u001b[39;00m \u001b[38;5;28mrange\u001b[39m(\u001b[38;5;28mlen\u001b[39m(prices_tensor) \u001b[38;5;241m-\u001b[39m window_size):\n\u001b[1;32m----> 6\u001b[0m   \u001b[43mX_train\u001b[49m\u001b[38;5;241;43m.\u001b[39;49m\u001b[43mappend\u001b[49m(prices_tensor[i:i\u001b[38;5;241m+\u001b[39mwindow_size])\n\u001b[0;32m      7\u001b[0m   y_train\u001b[38;5;241m.\u001b[39mappend(prices_tensor[i\u001b[38;5;241m+\u001b[39mwindow_size])\n\u001b[0;32m      9\u001b[0m   X_train \u001b[38;5;241m=\u001b[39m torch\u001b[38;5;241m.\u001b[39mstack(X_train)\n",
      "\u001b[1;31mAttributeError\u001b[0m: 'Tensor' object has no attribute 'append'"
     ]
    }
   ],
   "source": [
    "# Prepare sliding windows for training\n",
    "window_size = 60\n",
    "X_train, y_train = [], []\n",
    "\n",
    "for i in range(len(prices_tensor) - window_size):\n",
    "  X_train.append(prices_tensor[i:i+window_size])\n",
    "  y_train.append(prices_tensor[i+window_size])\n",
    "\n",
    "  X_train = torch.stack(X_train)\n",
    "  y_train = torch.stack(y_train)\n",
    "\n",
    "# Training loop\n",
    "epochs = 100\n",
    "for epoch in range(epochs):\n",
    "  model.train()\n",
    "  optimizer.zero_grad()\n",
    "  X_fft = fft.fft(X_train, dim=1) # Fourier Transform\n",
    "  output = model(X_fft.permute(0, 2, 1))\n",
    "  loss = criterion(output, y_train)\n",
    "  loss.backward()\n",
    "  optimizer.step()\n",
    "  if (epoch + 1) % 10 == 0:\n",
    "    print(f'Epoch [{epoch+1}/{epochs}], Loss: {loss.item():.4f}')"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "metadata": {},
   "outputs": [],
   "source": [
    "# Predicting prices over entire history using a sliding window\n",
    "model.eval()\n",
    "predicted_prices_history = []\n",
    "\n",
    "with torch.no_grad():\n",
    "  for i in range(len(prices_tensor) — window_size):\n",
    "  X_test_fft = fft.fft(prices_tensor[i:i+window_size].unsqueeze(0), dim=1)\n",
    "  predicted_prices = model(X_test_fft.permute(0, 2, 1)).squeeze().numpy()\n",
    "  predicted_prices_history.append(predicted_prices)\n",
    "\n",
    "predicted_prices_history = np.array(predicted_prices_history)\n",
    "# Plotting actual vs predicted prices for each stock\n",
    "plt.figure(figsize=(12, 8))\n",
    "for i, ticker in enumerate(tickers):\n",
    "plt.plot(data.index[window_size:], predicted_prices_history[:, i], label=f’Predicted — {ticker}’)\n",
    "plt.plot(data.index, data.values[:, i], label=f’Actual — {ticker}’)\n",
    "\n",
    "plt.xlabel(‘Date’)\n",
    "plt.ylabel(‘Stock Price’)\n",
    "plt.title(‘Actual vs. Predicted Stock Prices Over Time’)\n",
    "plt.legend()\n",
    "plt.show()"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "datahandler",
   "language": "python",
   "name": "python3"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.12.4"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 2
}
